"""
MetalX - Agent
接收消息格式：{"type":"command","message":"xxx", "args":"xxx"(可选)}
发送消息格式：
{"type":"error | response | notification | id","message":"xxx", 
"args":"xxx"(可选), "request_id":"xxx"（可选）}
API: send_notifications(dict)
API: connect()
API: disconnect()
配置项见init
"""

import socket
import json
import threading
import subprocess
import sys
import utils
import time
import configs


class Agent:
    def __init__(self, host, port=65432, max_retries=5, retry_interval=5):
        """
        初始化API客户端
        Args:
            host (str): 服务器地址
            port (number): 服务器端口，默认为65432
            max_retries (int): 最大重连次数，默认为5
            retry_interval (int): 重连间隔时间（秒），默认为5
        """
        self.host = host
        self.port = port
        self.client_socket = None
        self.running = False
        self.command_handlers = {
            "execute_command": self.handle_execute_command,
            "get_system_info": utils.get_system_info,
            "get_system_processes": utils.get_system_processes,
            "get_system_users": utils.get_system_users,
            "get_system_services": utils.get_system_services,
            "collect_server_status": utils.collect_server_status,
        }
        self.callbacks = {
            "on_connect": self.on_connect,
            "on_disconnect": self.on_disconnect,
            "on_command": self.on_command,
            "on_error": self.on_error,
        }
        self.max_retries = max_retries
        self.retry_interval = retry_interval
        self.retry_count = 0

    def send_response(self,request_id, result: dict):
        """发送命令执行结果给服务器"""
        message = {"type": "response", "message": result}
        self.send_message(message, request_id)


    def send_notification(self, content: dict):
        """主动发送通知给服务器"""
        message = {"type": "notification", "message": content}
        self.send_message(message)
   
    def send_error(self, content: dict):
        """主动发送错误给服务器"""
        message = {"type": "error", "message": content}
        self.send_message(message)
    
    
    def send_message(self, message: dict, request_id = None):
        """
        发送消息给服务器
        自定义消息类型
        """
        if not self.running or not self.client_socket:
            self.callbacks["on_error"]("无法发送消息: 未连接到服务器")
            return

        try:
            if(request_id):
                message.update({"request_id":request_id})
            self.client_socket.sendall(json.dumps(message).encode("utf-8"))
            
        except Exception as e:
            self.callbacks["on_error"](f"发送消息到服务器失败: {e}")
            self.cleanup()


        
    

    def connect(self):
        """
        连接到服务器
        Returns:
            bool:连接结果

        """
        self.client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            self.client_socket.connect((self.host, self.port))
            self.running = True
            receive_thread = threading.Thread(target=self.receive_messages, daemon=True)
            receive_thread.start()

            self.callbacks["on_connect"](f"已连接到服务器 {self.host}:{self.port}")
            return True
        except Exception as e:
            self.callbacks["on_error"](f"连接服务器失败: {e}")
            self.cleanup()
            return False

    def receive_messages(self):
        """
        接收服务器消息，函数会一直运行，直到捕捉到异常
        """
        try:
            while self.running:
                data = self.client_socket.recv(4096)
                if not data:
                    break

                try:
                    message = json.loads(data.decode("utf-8"))
                    # 将收到的信息交由处理函数处理
                    self.process_server_command(message)

                except json.JSONDecodeError as e:
                    self.callbacks["on_error"](f"无效的服务器消息: {data}")
        except ConnectionResetError:
            self.callbacks["on_disconnect"]("服务器断开连接")
        except Exception as e:
            self.callbacks["on_error"](f"接收消息时出错: {e}")
        finally:
            self.cleanup()

    def process_server_command(self, message):
        """
        处理来自服务器的消息
        Args:
            message(dict): 来自服务器的消息 {"type":"command","message":"xxx", "args":"xxx"(可选),request_id}
        """
        if "type" not in message:
            self.callbacks["on_error"]("无效的服务器消息: 缺少消息类型")
            return

        if message["type"] == "command":
            command = message.get("command")
            args = message.get("args", None)

            if command in self.command_handlers:
                self.callbacks["on_command"](command, args)
                try:
                    if args:
                        result = self.command_handlers[command](args)
                        self.send_response(message['request_id'],result)
                    else:
                        result = self.command_handlers[command]()
                        self.send_response(message['request_id'],result)
                except Exception as e:
                    self.send_notification(
                        {
                            "type": "error",
                            "message": f"执行命令 {command} 出错: {str(e)}",
                        }
                    )
                    self.callbacks["on_error"](f"执行命令 {command} 出错: {str(e)}")
            else:
                self.send_notification({"type": "error", "message": f"未知命令: {command}"})
                self.callbacks["on_error"](f"未知命令: {command}")
 

    def handle_execute_command(self, args) -> dict:
        """
        处理执行系统命令的命令
        Args:
            args(dict):要执行的命令{"command":"xxx"}

        Returns:
            dict: 运行结果：
            {
                "returncode": result.returncode,
                "stdout": result.stdout,
                "stderr": result.stderr,
            }
        """
        command = args.get("command")
        if not command:
            return {"type": "error", "message": "未提供命令"}

        try:
            result = subprocess.run(
                command,
                shell=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
            )
            return {
                "type": "result",
                "message": {
                    "returncode": result.returncode,
                    "stdout": result.stdout,
                    "stderr": result.stderr,
                },
            }
        except Exception as e:
            return {"type": "error", "message": f"执行命令出错: {str(e)}"}

    def disconnect(self):
        """断开与服务器的连接"""
        self.running = False
        if self.client_socket:
            try:
                self.client_socket.close()
            except Exception as e:
                self.callbacks["on_error"](f"清理资源时出错：{e}")
            self.client_socket = None

        self.callbacks["on_disconnect"]("客户端已断开连接")

    def cleanup(self):
        """清理资源,"""
        self.running = False
        if self.client_socket:
            try:
                self.client_socket.close()
            except Exception as e:
                self.callbacks["on_error"](f"清理资源时出错：{e}")
            self.client_socket = None

        self.callbacks["on_disconnect"]("客户端已断开连接")
        if self.retry_count < self.max_retries:
            self.callbacks["on_disconnect"](
                f"尝试第 {self.retry_count + 1} 次重连，{self.retry_interval} 秒后执行..."
            )
            time.sleep(self.retry_interval)
            if self.connect():
                self.retry_count = 0
            else:
                self.retry_count += 1

    def on_connect(self, info):
        print(f"[Agent]连接成功：{info}")

    def on_disconnect(self, reason):
        print(f"[Agent]断开连接: {reason}")

    def on_command(self, command, args):
        print(f"[Agent]收到命令: {command}, 参数: {args}")

    def on_error(self, error):
        print(f"[Agent]发生错误: {error}")
        self.send_error(f"[Agent]发生错误: {error}")


if __name__ == "__main__":
    # 可以从命令行参数获取服务器地址和端口
    host = configs.host
    port = configs.port

    if len(sys.argv) > 1:
        host = sys.argv[1]
    if len(sys.argv) > 2:
        port = int(sys.argv[2])

    client = Agent(host, port)

    if client.connect():
        try:
            # 保持主线程运行
            while True:
                time.sleep(1)
        except KeyboardInterrupt:
            client.disconnect()
